
class Point:

    def __init__(self,point):
        self.parentIndex=-1
        self.children=[]
        self.parent=None
        self.index=int(point[0])
        self.type=int(point[1])
        self.x=point[2]        
        self.y=point[3]
        self.z=point[4]
        self.ratio=point[5]
        self.parentIndex=int(point[6])
        self.children=[]
    def __str__(self):
        
        return str(self.index)+","+str(self.type)+","+str(self.x)+","+\
        str(self.y)+","+str(self.z)+","+str(self.ratio)+","+str(self.parentIndex)
        
class Edge:
    def __init__(self):
        self.data=[]
        pass
    def addPoint(self,point):
        data.append(point)
        
class NeuronTree:

    def __init__(self,filename):
        self.root=None
        self.edges=[]
        self.points=[]
        self.branchs=[]
        self.terminals=[]
        self.swc=[]
        with open(filename, "r") as f:
            text_lines = f.readlines()
            for line in text_lines:
                self.swc.append(list(map(float,(line.split()))))
            
        for c in self.swc:
            point = Point(c);
#            print(c) 
            self.points.append(point)
            if point.parentIndex==-1:
                self.root=point
            for p in reversed(self.points):
                if point.parentIndex==p.index:
                    point.parent=p
                    p.children.append(point)
                    if len(p.children)==2:
                        self.branchs.append(p)
#                    print( point)
                    break
        self.getEdges()
#         self.getGraph()
    def getEdges(self):
        self.edges.append([self.root]);
        for point in self.branchs:
            for i in range(len(point.children)):
                self.edges.append([point])
        for point in self.points:
            for i in range(len(self.edges)):
                edge=self.edges[i]
                lastp=edge[len(edge)-1]
                if point.parent==lastp and len(point.children)<=1:
                    edge.append(point)
                    break
        for edge in self.edges:
            if len(edge[len(edge)-1].children)==0:
                self.terminals.append(edge[len(edge)-1])
    
    def getGraph(self):
        edge=[self.root]
        self.getSubGraph(self.root,edge)
        pass
        
    def getSubGraph(self,point,edge):
        
        for child in point.children:
            subedge=edge[:]
            if len(child.children)==0:
                self.terminals.append(child);
            if len(child.children)==1:
                subedge.append(child);
            if len(child.children)>1:
                self.branchs.append(child);
            if len(child.children)==0 or len(child.children)>1:
                subedge.append(child)
                self.edges.append(subedge);
                subedge=[child]
            self.getSubGraph(child,subedge)
        pass
                    
                    
    def print(self):
        level='-';
        print(level,self.root.index)
        self.printChildren(self.root,level)
        
        
    def printChildren(self,point,level):
        level=level+'---'
        for child in point.children:
            print(level,child.index)
            self.printChildren(child,level)
            
    def printEdges(self):
        for edge in self.edges:
            print('==============')
            for p in edge:
                print (p.index);
        
if __name__ == "__main__":
    tree = NeuronTree("D:/project/FNTManager/build/install/data/swc/201625/005.swc")
    print("loaded tree")